import platform

import torch
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import transforms, datasets
from setting import batch_size, cpu_workers, data_dir, pin_memory
import cv2
import glob


class MyDataset(Dataset):

    def __init__(self, datas_path):
        super(MyDataset, self).__init__()
        # 读取数据
        x = []
        y = []
        labels = {
            'female': 0,
            'male': 1
        }

        for path in datas_path:
            img255 = cv2.imread(data_paths[0])
            x.append(img255.tolist())
            y.append(labels[path.split('\\' if platform.system() == 'Windows' else '/')[-2]])
        x = (torch.asarray(x) / 255.0).permute(0, 3, 1, 2)
        y = torch.asarray(y)
        self.data = (x, y)

    def __len__(self):
        x, y = self.data
        return len(x)

    def __getitem__(self, i):
        x, y = self.data
        return x[i], y[i]


data_paths = glob.glob(data_dir + '/real_data/*/*.jpg')
dataset = MyDataset(data_paths)
# train_dataset, test_dataset = random_split(dataset=dataset, lengths=(60, 10))

dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=True,
                        num_workers=cpu_workers, pin_memory=pin_memory)
# test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, drop_last=False,
#                              num_workers=cpu_workers, pin_memory=pin_memory)

if __name__ == '__main__':
    x, y = next(iter(dataloader))
    print(x.shape, y)
